# -*- coding:utf-8 -*-
__author__ = 'Dragon Sun'
__date__ = '2025-07-02 23:32:54'

import json
from typing import Any, Optional, Callable

from okx.websocket.WsUtils import initLoginParams
from websocket import WebSocket

from src.conf.conf import G_API配置, G_应用配置
from src.wss.ds_wss_okxv5_base import DSOKXv5WebSocketBase


class DS_OKXv5_WSS_私有频道类(DSOKXv5WebSocketBase):

    def __init__(self):
        super().__init__(url=G_API配置.wss_private)

        # 回调
        self.当前持仓变化回调: Optional[Callable[[DS_OKXv5_WSS_私有频道类], None]] = None
        self.收到USDT余额回调: Optional[Callable[[DS_OKXv5_WSS_私有频道类, float], None]] = None

    def _on_open(self, ws: WebSocket):
        super()._on_open(ws)

        self._登录()

    def _on_message(self, ws: WebSocket, message: Any):
        if self._为登录成功消息(message):
            self._订阅()

        if self.当前持仓变化回调 and self._为持仓变化消息(message):
            self.当前持仓变化回调(self)

        if self.收到USDT余额回调 and self._为余额消息(message):
            USDT余额 = self._获取币种余额(币种='USDT', message=message)
            if USDT余额:
                self.收到USDT余额回调(self, USDT余额)

    def _登录(self):
        payload = initLoginParams(
            useServerTime=False,
            apiKey=G_API配置.api_key,
            secretKey=G_API配置.secret_key,
            passphrase=G_API配置.passphrase
        )
        self.wss.send(payload)

    def _订阅(self):
        # 持仓
        param1 = {
            'channel': 'positions',
            'instType': G_应用配置.标的类型.value,
            'instFamily': G_应用配置.标的家族,
        }

        # 账户余额
        param2 = {
            'channel': 'balance_and_position',
        }

        # 订阅
        params = [param1, param2]
        self.订阅(params=params)

    @staticmethod
    def _为登录成功消息(message: str) -> bool:
        # noinspection PyBroadException
        try:
            """
                {
                    "event": "login",
                    "msg": "",
                    "code": "0",
                    "connId": "e88517c8"
                }            
            """
            d = json.loads(message)
            if (d['event'] == 'login') and (d['code'] == '0'):
                return True
            return False
        except Exception:
            return False

    @staticmethod
    def _为持仓变化消息(message: str) -> bool:
        """
            {
                "arg": {
                    "channel": "positions",
                    "instType": "SWAP",
                    "instFamily": "BTC-USDT",
                    "uid": "567121394413568"
                },
                "eventType": "snapshot",
                "curPage": 1,
                "lastPage": true,
                "data": [{
                    "adl": "1",
                    "availPos": "1",
                    "avgPx": "109111.5",
                    "baseBal": "",
                    "baseBorrowed": "",
                    "baseInterest": "",
                    "bePx": "109002.44302848577",
                    "bizRefId": "",
                    "bizRefType": "",
                    "cTime": "1751473604276",
                    "ccy": "USDT",
                    "clSpotInUseAmt": "",
                    "closeOrderAlgo": [{
                        "algoId": "2649874311530848256",
                        "closeFraction": "1.0",
                        "slTriggerPx": "110000",
                        "slTriggerPxType": "last",
                        "tpTriggerPx": "0",
                        "tpTriggerPxType": "last"
                    }],
                    "deltaBS": "",
                    "deltaPA": "",
                    "fee": "-0.5455575",
                    "fundingFee": "0",
                    "gammaBS": "",
                    "gammaPA": "",
                    "idxPx": "109564.0",
                    "imr": "",
                    "instId": "BTC-USDT-SWAP",
                    "instType": "SWAP",
                    "interest": "",
                    "last": "109504.7",
                    "lever": "10",
                    "liab": "",
                    "liabCcy": "",
                    "liqPenalty": "0",
                    "liqPx": "119484.86764559482",
                    "margin": "109.1115",
                    "markPx": "109504.8",
                    "maxSpotInUseAmt": "",
                    "mgnMode": "isolated",
                    "mgnRatio": "21.344269840226175",
                    "mmr": "4.380192",
                    "nonSettleAvgPx": "",
                    "notionalUsd": "1095.58457352",
                    "optVal": "",
                    "pTime": "1751474815213",
                    "pendingCloseOrdLiabVal": "",
                    "pnl": "0",
                    "pos": "1",
                    "posCcy": "",
                    "posId": "2649833903864811520",
                    "posSide": "short",
                    "quoteBal": "",
                    "quoteBorrowed": "",
                    "quoteInterest": "",
                    "realizedPnl": "-0.5455575",
                    "settledPnl": "",
                    "spotInUseAmt": "",
                    "spotInUseCcy": "",
                    "thetaBS": "",
                    "thetaPA": "",
                    "tradeId": "1770988974",
                    "uTime": "1751473604276",
                    "upl": "-3.933000000000029",
                    "uplLastPx": "-3.931999999999971",
                    "uplRatio": "-0.036045696374809",
                    "uplRatioLastPx": "-0.036036531438024",
                    "usdPx": "1.00049",
                    "vegaBS": "",
                    "vegaPA": ""
                }]
            }
        """
        # noinspection PyBroadException
        try:
            d: dict = json.loads(message)
            b1 = (d['arg']['channel'] == 'positions')
            b2 = (d['arg']['instType'] == G_应用配置.标的类型.value)
            b3 = (d['arg']['instFamily'] == G_应用配置.标的家族)
            b4 = (d['eventType'] == 'event_update')  # 持仓发生改变
            b5 = (len(d['data']) > 0)
            return b1 and b2 and b3 and b4 and b5
        except Exception:
            return False

    # noinspection SpellCheckingInspection
    @staticmethod
    def _为余额消息(message: str) -> bool:
        """
                    {
                        "arg": {
                            "channel": "balance_and_position",
                            "uid": "567121394413568"
                        },
                        "data": [{
                            "balData": [{
                                "cashBal": "0.0581838779610259",
                                "ccy": "BTC",
                                "uTime": "1751714869630"
                            }, {
                                "cashBal": "10",
                                "ccy": "LTC",
                                "uTime": "1615398325805"
                            }, {
                                "cashBal": "3000",
                                "ccy": "TUSD",
                                "uTime": "1615398325805"
                            }, {
                                "cashBal": "5",
                                "ccy": "ETH",
                                "uTime": "1615398325805"
                            }, {
                                "cashBal": "1000",
                                "ccy": "ADA",
                                "uTime": "1615398325805"
                            }, {
                                "cashBal": "83366.191404149915255",
                                "ccy": "USDT",
                                "uTime": "1751791664104"
                            }, {
                                "cashBal": "3000",
                                "ccy": "USDK",
                                "uTime": "1615398325805"
                            }, {
                                "cashBal": "500",
                                "ccy": "UNI",
                                "uTime": "1615398325805"
                            }, {
                                "cashBal": "100",
                                "ccy": "JFI",
                                "uTime": "1615398325805"
                            }, {
                                "cashBal": "10000",
                                "ccy": "TRX",
                                "uTime": "1615398325805"
                            }, {
                                "cashBal": "100",
                                "ccy": "OKB",
                                "uTime": "1615398325805"
                            }, {
                                "cashBal": "3000",
                                "ccy": "USDC",
                                "uTime": "1615398325805"
                            }, {
                                "cashBal": "3000",
                                "ccy": "PAX",
                                "uTime": "1615398325805"
                            }],
                            "eventType": "snapshot",
                            "pTime": "1751802305325",
                            "posData": [{
                                "avgPx": "107982.7",
                                "baseBal": "",
                                "ccy": "USDT",
                                "instId": "BTC-USDT-SWAP",
                                "instType": "SWAP",
                                "mgnMode": "isolated",
                                "nonSettleAvgPx": "",
                                "pos": "100",
                                "posCcy": "",
                                "posId": "2660495508046741504",
                                "posSide": "short",
                                "quoteBal": "",
                                "settledPnl": "",
                                "tradeId": "1795107576",
                                "uTime": "1751791664104"
                            }],
                            "trades": []
                        }]
                    }
                """
        # noinspection PyBroadException
        try:
            d: dict = json.loads(message)
            b1 = (d['arg']['channel'] == 'balance_and_position')
            b2 = (len(d['data'][0]['balData']) > 0)
            return b1 and b2
        except Exception:
            return False

    @staticmethod
    def _获取币种余额(币种: str, message: str) -> Optional[float]:
        # noinspection PyBroadException
        try:
            d: dict = json.loads(message)
            余额数组 = d['data'][0]['balData']
            for 余额项 in 余额数组:
                if 余额项['ccy'] == 币种:
                    余额 = float(余额项['cashBal'])
                    return 余额
            return None
        except Exception:
            return None


if __name__ == '__main__':
    ws_obj = DS_OKXv5_WSS_私有频道类()
    ws_obj.当前持仓变化回调 = lambda _: print('持仓变化')
    ws_obj.收到USDT余额回调 = lambda _, 余额: print(f'USDT余额 = {余额}')
    ws_obj.run()
